from typing import Text

from rasa.core.channels.channel import InputChannel, UserMessage, CollectingOutputChannel
from sanic import Blueprint, Sanic
from sanic.request import Request
from sanic.websocket import WebSocketProtocol, WebSocketConnection
import asyncio


class WebSocketChannel(InputChannel):
    def name(self) -> Text:
        return "websocket"

    def blueprint(self, on_new_message):
        websocket_blueprint = Blueprint('websocket_blueprint', __name__)

        @websocket_blueprint.websocket('/socket')
        async def websocket(request, ws):
            while True:
                message = await ws.recv()
                output_channel = CollectingOutputChannel()
                await on_new_message(UserMessage(message, output_channel, ws.id))
                responses = [message for _, message in output_channel.messages]
                for response in responses:
                    await ws.send(response["text"])

        return websocket_blueprint
